package lasec.util;

import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.mahout.cf.taste.common.NoSuchItemException;
import org.apache.mahout.cf.taste.common.NoSuchUserException;
import org.apache.mahout.cf.taste.common.Refreshable;
import org.apache.mahout.cf.taste.common.TasteException;
import org.apache.mahout.cf.taste.impl.common.FastIDSet;
import org.apache.mahout.cf.taste.impl.common.FullRunningAverage;
import org.apache.mahout.cf.taste.impl.common.LongPrimitiveIterator;
import org.apache.mahout.cf.taste.impl.common.RunningAverage;
import org.apache.mahout.cf.taste.impl.recommender.AbstractRecommender;
import org.apache.mahout.cf.taste.impl.recommender.TopItems;
import org.apache.mahout.cf.taste.model.DataModel;
import org.apache.mahout.cf.taste.model.Preference;
import org.apache.mahout.cf.taste.recommender.CandidateItemsStrategy;
import org.apache.mahout.cf.taste.recommender.IDRescorer;
import org.apache.mahout.cf.taste.recommender.RecommendedItem;
import org.apache.mahout.math.DenseMatrix;
import org.apache.mahout.math.Matrix;
import org.apache.mahout.math.SingularValueDecomposition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Preconditions;

public class SVDRecommender extends AbstractRecommender {
	private static final Logger log = LoggerFactory
			.getLogger(SVDRecommender.class);

	private double average;
	private double[][] doubleA;

	private Matrix A;

	private Map<String, Integer> rowLabels;
	private Map<String, Integer> columnLabels;

	public SVDRecommender(DataModel model, long userID) throws TasteException {
		this(model, getDefaultCandidateItemsStrategy(), userID);
	}

	public SVDRecommender(DataModel model,
			CandidateItemsStrategy defaultCandidateItemsStrategy, long userID)
			throws TasteException {
		// TODO Auto-generated constructor stub
		super(model, defaultCandidateItemsStrategy);

		int idx = 0;
		int numItems = model.getNumItems();

		LongPrimitiveIterator itemIterator = model.getItemIDs();
		rowLabels = new LinkedHashMap<String, Integer>(numItems);

		while (itemIterator.hasNext()) {
			rowLabels.put(itemIterator.next().toString(), idx++);
		}

		LongPrimitiveIterator userIterator = model.getUserIDs();
		idx = 0;
		int numUsers = model.getNumUsers();
		columnLabels = new LinkedHashMap<String, Integer>(numUsers);

		while (userIterator.hasNext()) {
			columnLabels.put(userIterator.next().toString(), idx++);
		}

		doubleA = new double[numItems][numUsers];
		A = new DenseMatrix(doubleA);
		A.setRowLabelBindings(rowLabels);
		A.setColumnLabelBindings(columnLabels);

		userIterator = model.getUserIDs();
		while (userIterator.hasNext()) {
			for (Preference pref : model.getPreferencesFromUser(userIterator
					.nextLong())) {
				A.set(String.valueOf(pref.getItemID()),
						String.valueOf(pref.getUserID()), pref.getValue());
			}
		}

		A = buildSVDMatrix(A);
		A.setRowLabelBindings(rowLabels);
		A.setColumnLabelBindings(columnLabels);

		average = getAveragePreference();
	}

	private Matrix buildSVDMatrix(Matrix matrix) {
		int k = 0;

		if (matrix.columnSize() > 10) {
			k = matrix.columnSize() / 3;
		} else {
			k = matrix.columnSize() % 10;
		}

		SingularValueDecomposition svd = new SingularValueDecomposition(matrix);

		Matrix U = svd.getU();
		Matrix Uk = U.viewPart(0, U.rowSize(), 0, k);

		Matrix rootSk = svd.getS().viewPart(0, k, 0, k);
		for (int i = 0; i < rootSk.columnSize(); i++) {
			rootSk.set(i, i, Math.sqrt(rootSk.get(i, i)));
		}

		Matrix V = svd.getV().transpose();
		Matrix Vk = V.viewPart(0, k, 0, V.columnSize());

		Matrix left = Uk.times(rootSk.transpose());
		Matrix right = rootSk.times(Vk);

		Matrix P = left.times(right);

		U = null;
		Uk = null;
		rootSk = null;
		V = null;
		Vk = null;
		left = null;
		right = null;

		return P;
	}

	private double getAveragePreference() throws TasteException {
		RunningAverage average = new FullRunningAverage();
		DataModel model = getDataModel();
		LongPrimitiveIterator it = model.getUserIDs();

		while (it.hasNext()) {
			for (Preference pref : model.getPreferencesFromUser(it.nextLong())) {
				average.addDatum(pref.getValue());
			}
		}

		return average.getAverage();
	}

	public String getKeysByValue(Map<String, Integer> map, int value) {
		Iterator<Entry<String, Integer>> iter = map.entrySet().iterator();
		String key = "";

		while (iter.hasNext()) {
			Map.Entry<String, Integer> entry = iter.next();
			if (entry.getValue().equals(value)) {
				key = entry.getKey();
			}
		}

		return key;
	}

	@Override
	public List<RecommendedItem> recommend(long userID, int howMany,
			IDRescorer rescorer) throws TasteException {
		// TODO Auto-generated method stub
		Preconditions.checkArgument(howMany >= 1, "howMany must be at least 1");
		log.debug("Recommending items for user ID '{}'", userID);

		int column = A.getColumnLabelBindings().get(String.valueOf(userID));

		FastIDSet possibleItemIDs = new FastIDSet();
		for (int i = 0; i < A.rowSize(); i++) {
			if (A.get(i, column) <= 0) {
				possibleItemIDs.add(Long.parseLong(getKeysByValue(
						A.getRowLabelBindings(), i)));
			}
		}

		TopItems.Estimator<Long> estimator = new Estimator(userID);

		List<RecommendedItem> topItems = TopItems.getTopItems(howMany,
				possibleItemIDs.iterator(), rescorer, estimator);

		log.debug("Recommendations are: {}", topItems);

		A = null;
		doubleA = null;
		columnLabels = null;
		rowLabels = null;
		
		return topItems;
	}

	@Override
	public float estimatePreference(long userID, long itemID)
			throws TasteException {
		// TODO Auto-generated method stub
		Integer useridx = columnLabels.get(String.valueOf(userID));
		if (useridx == null) {
			throw new NoSuchUserException();
		}

		Integer itemidx = rowLabels.get(String.valueOf(itemID));
		if (itemidx == null) {
			throw new NoSuchItemException();
		}

		return predictRating(useridx, itemidx);
	}

	private float predictRating(int user, int item) {
		float entry = (float) A.get(item, user);
		float prediction = (float) (entry + average);

		if (prediction <= 1) {
			return 1.0f;
		} else if (prediction > 5) {
			return 5.0f;
		} else {
			return prediction;
		}
	}

	@Override
	public void refresh(Collection<Refreshable> alreadyRefreshed) {
		// TODO Auto-generated method stub

	}

	private final class Estimator implements TopItems.Estimator<Long> {
		private final long theUserID;

		private Estimator(long theUserID) {
			this.theUserID = theUserID;
		}

		@Override
		public double estimate(Long itemID) throws TasteException {
			// TODO Auto-generated method stub
			return estimatePreference(theUserID, itemID);
		}
	}
}